import torch
import os
from torchvision import datasets


class muti_targets(datasets.ImageFolder):
    def __getitem__(self, index):
        original_tuple = super(muti_targets, self).__getitem__(index)

        path = self.imgs[index][0]

        tuple_with_path = (original_tuple + (path,))

        return tuple_with_path

    def index2classlist(self):
        return self._find_classes_(self.root)

    def _find_classes_(self, dir):
        """
        list : index of list coresponding to classname
        """
        num_class = len(os.listdir(dir)[0])
        muti_class_list = []
        for i in range(num_class):
            classes_list = [d[i] for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))]
            set_classes = list(set(classes_list))
            set_classes.sort()
            muti_class_list.append(set_classes)

        return muti_class_list